import io
import numpy as np
import streamlit as st
from PIL import Image
from . import map_tab



from exg import exg_stress_mask, overlay_mask
from yolo_demo import run_yolo_on_image

st.set_page_config(page_title="PestWatch Lite", layout="wide")
st.title("🌾 PestWatch Lite — Day 1 演示")


with st.sidebar:
    st.markdown("**小贴士**：今天先验证端到端流程，明天再换成你自己的虫害/病斑模型。\n\n"
                "- 左侧切换标签：YOLO占位推理 / ExG植被指数\n"
                "- 上传RGB图片（JPG/PNG）\n"
                "- ExG用于估计低绿度（可能应激区），仅作近似演示")


tab1, tab2, tab3 = st.tabs(["🔍 YOLO 占位推理", "🌱 ExG 植被指数", "🗺️ 定位与地图"])

with tab1:
    st.subheader("YOLO 占位推理（验证前后端管线）")
    img_file = st.file_uploader("上传一张RGB图片（JPG/PNG）", type=["jpg","jpeg","png"], key="yolo_upl")
    colL, colR = st.columns([1,1])
    if img_file is not None:
        pil_img = Image.open(io.BytesIO(img_file.read())).convert("RGB")
        with st.spinner("运行模型中…（首次会自动下载yolov8n权重）"):
            try:
                plotted_rgb, dets = run_yolo_on_image(pil_img)
            except Exception as e:
                st.error(f"推理失败：{e}")
            else:
                with colL:
                    st.image(pil_img, caption="原图", use_column_width=True)
                with colR:
                    st.image(plotted_rgb, caption="检测可视化（COCO通用类别）", use_column_width=True)
                st.markdown("**检测结果（前5）**")
                for d in dets[:5]:
                    st.write(f"- {d['name']}  置信度: {d['conf']:.2f}  框: {list(map(lambda x: round(x,1), d['box']))}")
                if not dets:
                    st.info("未检测到目标。注意：这是通用模型，不包含虫害特定类别。")
    else:
        st.info("请先上传图片。")

with tab2:
    st.subheader("ExG 植被指数（RGB近似）")
    img_file2 = st.file_uploader("上传一张RGB图片（JPG/PNG）", type=["jpg","jpeg","png"], key="exg_upl")
    if img_file2 is not None:
        pil_img2 = Image.open(io.BytesIO(img_file2.read())).convert("RGB")
        rgb = np.array(pil_img2)
        with st.spinner("计算ExG与低绿度掩膜…"):
            mask, exg_u8, thr = exg_stress_mask(rgb)
            covered = float(mask.sum()) / mask.size * 100.0
            overlay = overlay_mask(rgb, mask, alpha=0.45)

        col1, col2, col3 = st.columns([1,1,1])
        with col1:
            st.image(rgb, caption="原图", use_column_width=True)
        with col2:
            st.image(exg_u8, caption=f"ExG（灰度）- Otsu阈值: {thr}", use_column_width=True, clamp=True)
        with col3:
            st.image(overlay, caption=f"低绿度区（红色）≈ {covered:.1f}%", use_column_width=True)


    else:
        st.info("请先上传图片。")
with tab3:
    map_tab.render()